#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Dec 26 23:20:25 2022

@author: matthewroberts
"""

import os
import sys
import numpy as np
import glob

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patheffects as path_effects
import matplotlib.gridspec as gridspec
import matplotlib.colors as colors
import matplotlib.cm as cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap

import cartopy.crs as ccrs
import pyproj

import pyart
import datetime as dt
import wrf
from netCDF4 import Dataset
import pandas as pd
from pyproj import Geod
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import cartopy.io.shapereader as shpreader
from cartopy.io.img_tiles import GoogleTiles
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter

import warnings
warnings.filterwarnings('ignore')

start_time = dt.datetime.now()
print('\nProgram Started: {0}'.format(start_time))
print('===================')

def relax_zone_remover (input, sr):
    # remove extra points in fire grid
    output = input
    for _ in range(sr):
        output = np.delete(output, -1, 0)
        output = np.delete(output, -1, 1)
    return output

#########################
# Inputs
#########################

firename = 'bear_fire' # for naming directories/figures
extent = [-121.5,-120.9, 39.48,39.9] #bear fire
idx0 = 0
idx1 = -1

# firename = 'caldor_fire' # for naming directories/figures
# extent = [-120.61, -120.28, 38.52, 38.76] #caldor fire
# idx0 = 0
# idx1 = -2

keyword = '_FUELx8'

mainpath = '/Users/matthewroberts/Documents/Projects/LEAPHI'
filepath = mainpath+'/data/'+firename+'/'+keyword+'/' #data files location
savepath = mainpath+'/plots/'+firename+'/' #where to save figure


# get last file and perimeter
wrf_file1 = Dataset(sorted(glob.glob(filepath+'/wrfout_d01*'))[0],mode='r') # for initial fuel map
wrf_file2 = Dataset(sorted(glob.glob(filepath+'/wrfout_d02*'))[0],mode='r') #for final perimeter
wrf_file3 = Dataset(sorted(glob.glob(filepath+'/wrfout_d03*'))[0],mode='r') #for final perimeter

# Standard lat/lon grid
terrain1 = wrf.getvar(wrf_file1, "HGT", timeidx=-1)
lat1 = wrf.getvar(wrf_file1, "XLAT", timeidx=-1)
lon1 = wrf.getvar(wrf_file1, "XLONG", timeidx=-1)

lat2 = wrf.getvar(wrf_file2, "XLAT", timeidx=-1)
lon2 = wrf.getvar(wrf_file2, "XLONG", timeidx=-1)

lat3 = wrf.getvar(wrf_file3, "XLAT", timeidx=-1)
lon3 = wrf.getvar(wrf_file3, "XLONG", timeidx=-1)

#########################
# Plotting
#########################
print('Plotting...', end="", flush=True)
fig = plt.figure(1,figsize=(9,9))
ax1 = fig.add_subplot(111, projection=ccrs.PlateCarree()) #rows,cols (1)
extent = [lon1[0,0], lon1[0,-1], lat1[1,0], lat1[-1,0]]
ax1.set_extent(extent, crs=ccrs.PlateCarree())

# Draw the fuelmap
cmap1 = cm.get_cmap('terrain',lut=30)

# fuelcopy[lfn1 > 0] = np.nan
# fuelcopy[fuel.data == 165] = 999

# val = 19
# fuelcopy = np.ma.masked_less(fuelcopy, val)
# fuelcopy = np.ma.masked_greater(fuelcopy, val)

terrainplot = ax1.pcolormesh(lon1, lat1, terrain1, cmap=cmap1, vmin=-850, vmax=2800, transform=ccrs.PlateCarree())
# Add reference cities
text1 = ax1.text(-119.810535, 39.521532,'Reno', horizontalalignment='center', fontsize=11, fontweight='bold', color='k', alpha=.7, zorder=10)
text2 = ax1.text(-121.475419, 38.565623, 'Sacramento', horizontalalignment='center', fontsize=11, fontweight='bold', color='k', alpha=.7, zorder=10)
# text1.set_path_effects([path_effects.Stroke(linewidth=1, foreground='k', alpha=.6),
#                        path_effects.Normal()])
# text2.set_path_effects([path_effects.Stroke(linewidth=1, foreground='k', alpha=.6),
#                        path_effects.Normal()])
                       
bds = ax1.plot(lon2[0,:], lat2[0,:], linewidth=3, color='k', zorder=11)
bds = ax1.plot(lon2[-1,:], lat2[-1,:], linewidth=3, color='k', zorder=11)
bds = ax1.plot(lon2[:,0], lat2[:,0], linewidth=3, color='k', zorder=11)
bds = ax1.plot(lon2[:,-1], lat2[:,-1], linewidth=3, color='k', zorder=11)

bds = ax1.plot(lon3[0,:], lat3[0,:], linewidth=3, color='k', zorder=11)
bds = ax1.plot(lon3[-1,:], lat3[-1,:], linewidth=3, color='k', zorder=11)
bds = ax1.plot(lon3[:,0], lat3[:,0], linewidth=3, color='k', zorder=11)
bds = ax1.plot(lon3[:,-1], lat3[:,-1], linewidth=3, color='k', zorder=11)

ax1.set_xticks(np.linspace(extent[0],extent[1],7),crs=ccrs.PlateCarree()) # set longitude indicators
ax1.set_yticks(np.linspace(extent[2],extent[3],8)[1:],crs=ccrs.PlateCarree()) # set latitude indicators
lon_formatter = LongitudeFormatter(number_format='0.2f',degree_symbol=u'\N{DEGREE SIGN}',dateline_direction_label=True) # format lons
lat_formatter = LatitudeFormatter(number_format='0.2f',degree_symbol=u'\N{DEGREE SIGN}') # format lats
ax1.xaxis.set_major_formatter(lon_formatter) # set lons
ax1.yaxis.set_major_formatter(lat_formatter) # set lats
ax1.xaxis.set_tick_params(width=2,length=5,direction='in',labelsize=12, zorder=12, rotation=20, bottom=True, top=True)
ax1.yaxis.set_tick_params(width=2,length=5,direction='in',labelsize=12, zorder=12, left=True, right=True)

#Download county shapefiles
reader = shpreader.Reader('/Users/matthewroberts/Documents/Data/shapefiles/countyl010g_shp/countyl010g.shp')
counties = list(reader.geometries())
COUNTIES = cfeature.ShapelyFeature(counties, ccrs.PlateCarree())
ax1.add_feature(COUNTIES, facecolor='none', edgecolor='k', linewidth=1, alpha=.4, zorder=10)

ax1.outline_patch.set_linewidth(2)
ax1.outline_patch.set_zorder(99)

ax1.set_xlabel("Longitude", fontsize=12)
ax1.set_ylabel("Latitude", fontsize=12)

# color_idx = np.arange(len(fuelkey))
cbar = plt.colorbar(terrainplot, ax=ax1, shrink=.55)#, ticks=ticks)
cbar.set_label('Terrain Height [m MSL]')
# cbar.set_ticklabels(fuelkey)
    
# plt.colorbar(fuelplot)

plt.savefig(savepath+'/domainmap.png',bbox_inches='tight',dpi=900)
# plt.show()
# sys.exit()








